
#include "NodeTable.hpp"
#include "Common.hpp"
#include <algorithm>
#include <random>

KnowAddress::KnowAddress(muduo::net::InetAddress const& _addr,muduo::net::InetAddress const& _src):
            addr(_addr),src(_src),attempts(0),lastAttempt(muduo::Timestamp::now()),bucketType(bucketTypeNew)
{}

int KnowAddress::addBucketRef(int bucketIdx)
{
    if(std::find(buckets.begin(),buckets.end(),bucketIdx)!=buckets.end()) return -1;
    buckets.push_back(bucketIdx);
    return buckets.size();
}

int KnowAddress::removeBucketRef(int bucketIdx)
{
    std::vector<int> res;
    for(auto i=buckets.begin();i!=buckets.end();i++)
    {
        if(*i!=bucketIdx) res.push_back(*i);
    }
    if(buckets.size()!=res.size()+1) return -1;
    buckets.swap(res);
    return buckets.size();
}

bool KnowAddress::isBad()
{
    // Is Old --> good
    if(bucketType==bucketTypeOld) return false;

    muduo::Timestamp now = muduo::Timestamp::now();

    // Has been attempted in the last minute --> good
    if(muduo::timeDifference(now, lastAttempt) < 60) return false;

    
    if(muduo::timeDifference(now,lastAttempt) > kNumMissingDays*24*60) 
        return true;

    // Never succeeded?
    if(lastSuccess.microSecondsSinceEpoch()==0 && attempts > kNumRetries)
        return true;
    
    if(muduo::timeDifference(now, lastSuccess) > kMinBadDays*24*60 && attempts >= kMaxFailures)
        return true;
    
    return false;
}



NodeTable::NodeTable(std::string filePath, bool routabilityStrict,std::shared_ptr<spdlog::logger> logger):
        filePath_(filePath),rotabilityStrict_(routabilityStrict), logger_(logger)
{
    for(auto i=bucketsNew_.begin();i!=bucketsNew_.end();i++)
    {
        *i=std::unordered_map<std::string, std::shared_ptr<KnowAddress>>();
    }
    for(auto i=bucketsOld_.begin();i!=bucketsOld_.end();i++)
    {
        *i=std::unordered_map<std::string, std::shared_ptr<KnowAddress>>();
    }
}

void NodeTable::addOurAddress(muduo::net::InetAddress const&  addr)
{
    std::unique_lock<std::mutex> lock(mtx_);
    logger_->info("Add our address to book addr: {}", addr.toIpPort());
    ourAddr_.insert(addr.toIpPort());
}

bool NodeTable::isOurAddress(muduo::net::InetAddress const&  addr)
{
    std::unique_lock<std::mutex> lock(mtx_);
    return ourAddr_.count(addr.toIpPort()) == 1;
}

void NodeTable::addAddress(muduo::net::InetAddress const& addr, muduo::net::InetAddress const& src)
{
    std::unique_lock<std::mutex> lock(mtx_);

    if(ourAddr_.count(addr.toIpPort())==1)
    {
        logger_->warn("Cannot add ourselves with address {}",addr.toIpPort());
        return;
    }
    auto k = addrLookup_.find(addr.toIpPort());
    if(k != addrLookup_.end())
    {
        // the addr already in addrLookup_
        if(k->second->isOld()) return;
        if(k->second->buckets.size() == kMaxNewBucketsPerAddress) return;
    }
    else 
    {
        addrLookup_[addr.toIpPort()] = std::make_shared<KnowAddress>(addr, src);
    }

    int bucket = calcNewBucket(addr, src);
    bool added = addToNewBucket(addrLookup_[addr.toIpPort()], bucket);
    if(!added)
    {
        logger_->info("Can't add new address, addr book is full address:{} total: {}",addr.toIpPort(),sizeWithoutMutex());
    }
    else
    {
        logger_->info("Added new address: {} total: {}",addr.toIpPort(),sizeWithoutMutex());
    }
}

std::vector<std::shared_ptr<KnowAddress>> NodeTable::listOfKnownAddress()
{
    std::unique_lock<std::mutex> lock(mtx_);
    std::vector<std::shared_ptr<KnowAddress> > res;
    for(auto i=addrLookup_.begin();i!=addrLookup_.end();i++)
    {
        res.push_back(std::make_shared<KnowAddress>(*(i->second)));
    }
    return res;
}

std::vector<muduo::net::InetAddress> NodeTable::getSelection()
{
    std::unique_lock<std::mutex> lock(mtx_);
    std::vector<muduo::net::InetAddress> allAddr;
    if(sizeWithoutMutex() == 0) return allAddr;
    
    allAddr.reserve(sizeWithoutMutex());
    for(auto i=addrLookup_.begin();i!=addrLookup_.end();i++)
    {
        allAddr.push_back(i->second->addr);
    }

    int numAddresses = std::max(std::min(kMinGetSelection,(int)allAddr.size()),
       int(allAddr.size()*kGetSelectionPercent/100));
    numAddresses=std::min(kMaxGetSelection, numAddresses);
    std::mt19937 gen(1729);
    for(int i=0;i<numAddresses;i++)
    {
        std::uniform_int_distribution<> dis(0,allAddr.size()-i);
        int j=dis(gen)+i;
        std::swap(allAddr[i],allAddr[j]);
    }
    return std::vector<muduo::net::InetAddress>(allAddr.begin(),allAddr.begin()+numAddresses);
}

std::unordered_map<std::string, std::shared_ptr<KnowAddress>> & NodeTable::getBucket(int bucketType, int bucketIdx)
{
    if(bucketType==bucketTypeNew)
    {
        return bucketsNew_[bucketIdx];
    }
    else return bucketsOld_[bucketIdx];
}


bool NodeTable::addToNewBucket(std::shared_ptr<KnowAddress> & ka, int bucketIdx)
{
    if(ka->isOld())
    {
        logger_->error("Cannot add address already in old bucket to a new bucket: {}",ka->addr.toIpPort());
        return false;
    }

    std::string addrStr = ka->addr.toIpPort();
    auto bucket = getBucket(bucketTypeNew, bucketIdx);
    // Already exists?
    if(bucket.find(addrStr) != bucket.end()) return true;
    
    // Enforce max address.
    if(bucket.size() > kNewBucketSize)
    {
        logger_->info("new bucket is full, expiring new");
        expireNew(bucketIdx);
    }

    // Add to bucket.
    bucket[addrStr]= ka;
    if(ka->addBucketRef(bucketIdx)==1)
        nNew++;

    addrLookup_[addrStr] = ka;
    return true;
}    


bool NodeTable::addToOldBucket(std::shared_ptr<KnowAddress>& ka, int bucketIdx)
{
    // Sanity check
    if(ka->isNew())
    {
        logger_->error("Cannot add new address to old bucket: ", ka->addr.toIpPort());
        return false;
    }
    if(!ka->buckets.empty())
    {
        logger_->error("Cannot add already old address to another old bucket: ", ka->addr.toIpPort());
        return false;
    }

    auto addrStr = ka->addr.toIpPort();
    auto bucket = getBucket(bucketTypeOld, bucketIdx);

    // Already exists?
    if(bucket.find(addrStr) != bucket.end()) return true;

    // Enforce max addresses.
    if(bucket.size() > kOldBucketSize) return false;

    // Add to bucket
    bucket[addrStr] = ka;
    if(ka->addBucketRef(bucketIdx)==1) nOld++;

    addrLookup_[addrStr] = ka;

    return true;
}

void NodeTable::removeFromBucket(std::shared_ptr<KnowAddress>& ka, int bucketType, int bucketIdx)
{
    if(ka->bucketType != bucketType)
    {
        logger_->error("removeFromBucket: Bucket type mismatch: {}",ka->addr.toIpPort());
        return;
    }
    auto bucket = getBucket(bucketType, bucketIdx);
    bucket.erase(ka->addr.toIpPort());
    if (ka->removeBucketRef(bucketIdx) == 0)
    {
        if(bucketType == bucketTypeNew) nNew--;
        else nOld--;
        addrLookup_.erase(ka->addr.toIpPort());
    }
}

void NodeTable::removeFromAllBucket(std::shared_ptr<KnowAddress>& ka)
{
    for(int i=0;i< ka->buckets.size();i++)
    {
        auto bucket=getBucket(ka->bucketType,ka->buckets[i]);
        bucket.erase(ka->addr.toIpPort());
    }
    ka->buckets.clear();
    if(ka->bucketType == bucketTypeNew) nNew--;
    else nOld--;
    addrLookup_.erase(ka->addr.toIpPort());
}

std::shared_ptr<KnowAddress> NodeTable::pickOldest(int bucketType, int bucketIdx)
{
    auto bucket = getBucket(bucketType, bucketIdx);
    std::shared_ptr<KnowAddress> oldest;
    for(auto i=bucket.begin();i!=bucket.end();i++)
    {
        if(oldest.use_count()==0 || i->second->lastAttempt < oldest->lastAttempt)
            oldest=i->second;
    }
    return oldest;
}

void NodeTable::moveToOld(std::shared_ptr<KnowAddress> &ka)
{
    if(ka->isOld())
    {
        logger_->error("Cannot promote address that is already old {}", ka->addr.toIpPort());
        return;
    }
    if(ka->buckets.size()==0)
    {
        logger_->error("Cannot promote address that isn't in any new buckets {}", ka->addr.toIpPort());
        return;
    }
    // Remeber one of the buckets in which ka is in
    int freedBucket = ka->buckets[0];
    // Remove from all (new) buckets
    removeFromAllBucket(ka);
    // It is officially old now
    ka->bucketType = bucketTypeNew;

    // Try to add it to its oldBucket destination
    auto oldBucketIdx = calcOldBucket(ka->addr);
    bool added = addToOldBucket(ka, oldBucketIdx);

    if(!added)
    {
        // No room, must evict something
        auto oldest=pickOldest(bucketTypeOld, oldBucketIdx);
        removeFromBucket(oldest, bucketTypeOld,oldBucketIdx);

        //find new bucket to put the oldest in
        auto newBucketIdx = calcNewBucket(oldest->addr, oldest->src);
        added = addToNewBucket(oldest, newBucketIdx);

        // No space in newBucket either, just put it in freedBucket from above
        if(!added)
        {
            added= addToNewBucket(oldest, freedBucket);
            if(!added)
            {
                logger_->error("Could not migrate oldest {} to freeBucket {}",oldest->addr.toIpPort(), freedBucket);
            }
        }
        added = addToOldBucket(ka, oldBucketIdx);
        if(!added)
        {
            logger_->error("Could not re-add ka {} to oldBucketIdx {}",ka->addr.toIpPort(),oldBucketIdx);
        }
    }
}


void NodeTable::expireNew(int bucketIdx)
{
    for(auto i=bucketsNew_[bucketIdx].begin();i!=bucketsNew_[bucketIdx].end();i++)
    {
        if(i->second->isBad())
        {
            logger_->info("expiring bad address {}", i->first);
            removeFromBucket(i->second, bucketTypeNew, bucketIdx);
            return;
        }
    }
    auto oldest = pickOldest(bucketTypeNew, bucketIdx);
    removeFromBucket(oldest, bucketTypeNew, bucketIdx);
}

bool NodeTable::hasAddress(muduo::net::InetAddress const& addr)
{
    std::unique_lock<std::mutex> lock(mtx_);
    return addrLookup_.find(addr.toIpPort())!=addrLookup_.end();
}

bool NodeTable::needMoreAddress()
{
    return sizeWithMutex() < kNeedAddressThreshold;
}


void markGood(muduo::net::InetAddress const& addr);
    void markAttemp(muduo::net::InetAddress const& addr);
    void markBad(muduo::net::InetAddress const& addr); 
// IsGood returns true if peer was ever marked as good and haven't
// done anything wrong since then.
bool NodeTable::isGood(muduo::net::InetAddress const& addr)
{
    std::unique_lock<std::mutex> lock(mtx_);
    return addrLookup_[addr.toIpPort()]->isOld();
}

void NodeTable::markGood(muduo::net::InetAddress const& addr)
{
    std::unique_lock<std::mutex> lock(mtx_);
    if(addrLookup_.find(addr.toIpPort()) == addrLookup_.end())  return;

    auto ka=addrLookup_[addr.toIpPort()];
    ka->markGood();

    if(ka->isNew()) moveToOld(ka);
}

void NodeTable::markAttemp(muduo::net::InetAddress const& addr)
{
    std::unique_lock<std::mutex> lock(mtx_);
    if(addrLookup_.find(addr.toIpPort()) == addrLookup_.end()) return;
    addrLookup_[addr.toIpPort()]->markAttemp();
}

void NodeTable::markBad(muduo::net::InetAddress const& addr)
{
    removeAddress(addr);
}

int NodeTable::calcNewBucket(const muduo::net::InetAddress &addr, const muduo::net::InetAddress& src)
{
    std::hash<std::string> h;
    return h(addr.toIpPort()+src.toIpPort())%kNewBucketCount;
}

int NodeTable::calcOldBucket(const muduo::net::InetAddress &addr)
{
    std::hash<std::string> h;
    return  h(addr.toIpPort())%kOldBucketCount;
}